/*
 * Copyright (c) 2021 The red-star Project
 *
 * Licensed under the Apache License, version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at:
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.inyourcode.core.db.iml;

import com.inyourcode.core.db.api.IDBResource;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mongodb.client.model.UpdateOptions;
import org.bson.Document;
import org.bson.types.Binary;
import org.springframework.data.mongodb.core.MongoTemplate;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @author JackLei
 */
public class MongoResource implements IDBResource {
    /** 唯一键，统一用这个名称*/
    private static final String PRIMARY_KEY_NAME = "pid";
    private MongoTemplate mongoTemplate;
    private String collectionName;
    private MongoCollection<Document> collection;
    private static UpdateOptions updateOptions;

    public MongoResource(MongoTemplate mongoTemplate, String collectionName) {
        this.mongoTemplate = mongoTemplate;
        this.collectionName = collectionName;
        this.collection = mongoTemplate.getCollection(collectionName);
        this.updateOptions = new UpdateOptions();
        this.updateOptions.upsert(true);
    }

    @Override
    public Map<String, byte[]> load(String key) {
        Document query = new Document();
        query.append(PRIMARY_KEY_NAME, key);

        FindIterable<Document> documents = this.collection.find(query);
        if (documents == null) {
            return null;
        }

        MongoCursor<Document> iterator = documents.iterator();
        if (iterator.hasNext()) {
            Document document = iterator.next();
            if (document == null) {
                return null;
            }

            Map<String, byte[]> dataMap = new HashMap<>();
            Set<Map.Entry<String, Object>> entries = document.entrySet();
            for (Map.Entry<String, Object> entry : entries) {
                String tableName = entry.getKey();
                Object tableData = entry.getValue();
                if (!Binary.class.isAssignableFrom(tableData.getClass())) {
                    continue;
                }

                dataMap.put(tableName, ((Binary)tableData).getData());
            }

            return dataMap;
        }

        return null;
    }

    @Override
    public void save(String key, Map<String, byte[]> updateDatas) {
        Document updateDoc = new Document();
        for (Map.Entry<String, byte[]> entry : updateDatas.entrySet()) {
            String subKey = entry.getKey();
            byte[] data = entry.getValue();
            updateDoc.append(subKey, data);
        }
        Document updateKey = new Document(PRIMARY_KEY_NAME, key);
        Document setDocument = new Document("$set", updateDoc);

        this.collection.updateOne(updateKey, setDocument, updateOptions);
    }
}
